#!/bin/python3
# 本代码实现hardware events

from bcc import BPF, PerfHWConfig, PerfType
from ctypes import *
import argparse
import signal
import sys

fun_start = 'fun_start'     # 函数进入时触发的bpf函数名
fun_end = 'fun_end'         # 函数退出时触发的bpf函数名
fun_output = 'output'       # BPF_PERF_OUTPUT结构名
tfun_bpf_hash_name = 'tfun_hash'  # bpf哈希表名称

base_bpf = '''
#include <uapi/linux/ptrace.h>
#include <uapi/linux/bpf_perf_event.h>

struct perf_delta{
    u64 time_delta; //时间
};

#define TIME_ID 0
BPF_PERF_OUTPUT(output);
BPF_PERCPU_ARRAY(data, u64);
BPF_HASH('''+tfun_bpf_hash_name+''',u32);
'''

tfun_bpf_fname_start = 'fun_start'
tfun_bpf_fname_end = 'fun_end'
tfun_bpf_fname_cpu_cycles = 'on_cpu_cycles'
tfun_bpf_fname_instructions = 'on_instructions'
tfun_bpf_fname_cache_misses = 'on_cache_misses'
tfun_bpf_fname_cache_references = 'on_cache_references'
tfun_bpf_fname_branch_misses = 'on_branch_misses'
tfun_bpf_fname_branch_instructions = 'on_branch_instructions'
tfun_bpf_fname_bus_cycles = 'on_bus_cycles'


class tfun_bpf_hash:
    CPU_CYCLES = 0
    INSTRUCTIONS = 1
    CACHE_REFERENCES = 2
    CACHE_MISSES = 3
    BRANCH_INSTRUCTIONS = 4
    BRANCH_MISSES = 5
    BUS_CYCLES = 0


def tfun_bpf_template(fname, hash_key):
    #start='BPF_HASH('+fname[3:]+',u32);\nint '
    tmp = '\nint '+fname+'''(struct bpf_perf_event_data *ctx) {
    u32 key='''+str(hash_key)+''';
    u32 timeid=TIME_ID;
    u64 *kptr=data.lookup(&timeid);
    if(kptr && *kptr!=0){
            tfun_hash.increment(key,ctx->sample_period);
    }
    return 0;
}
    '''
    return tmp


tfun_bpf_cpu_cycles = tfun_bpf_template(
    tfun_bpf_fname_cpu_cycles, tfun_bpf_hash.CPU_CYCLES)
tfun_bpf_instructions = tfun_bpf_template(
    tfun_bpf_fname_instructions, tfun_bpf_hash.INSTRUCTIONS)
tfun_bpf_cache_misses = tfun_bpf_template(
    tfun_bpf_fname_cache_misses, tfun_bpf_hash.CACHE_MISSES)
tfun_bpf_cache_references = tfun_bpf_template(
    tfun_bpf_fname_cache_references, tfun_bpf_hash.CACHE_REFERENCES)
tfun_bpf_branch_misses = tfun_bpf_template(
    tfun_bpf_fname_branch_misses, tfun_bpf_hash.BRANCH_MISSES)
tfun_bpf_branch_instructions = tfun_bpf_template(
    tfun_bpf_fname_branch_instructions, tfun_bpf_hash.BRANCH_INSTRUCTIONS)
tfun_bpf_bus_cycles = tfun_bpf_template(
    tfun_bpf_fname_bus_cycles, tfun_bpf_hash.BUS_CYCLES)


def get_tfun_bpf_start(dis_time=False):
    fun_text = '\nvoid '+tfun_bpf_fname_start+'''(struct pt_regs *ctx)
{
    u32 time=TIME_ID;
    '''
    if dis_time:
        fun_text += 'u64 time_start =1;\n\t'
    else:
        fun_text += 'u64 time_start=bpf_ktime_get_ns();\n\t'
    fun_text += '''u64 *kptr = data.lookup(&time);
    if (kptr) {
        data.update(&time, &time_start);
    } else {
        data.insert(&time, &time_start);
    }
}
    '''
    return fun_text


def get_tfun_bpf_end(dis_time=False):
    fun_text = '\nvoid '+tfun_bpf_fname_end+'''(struct pt_regs *ctx)
{
    u32 time = TIME_ID;
    u64 timeid=0;
    '''
    if dis_time:
        fun_text += '''
    u64 *kptr = data.lookup(&time);
    if (kptr) {
        data.update(&time, &timeid);
    } else {
        data.insert(&time, &timeid);
        return;
    }
}   
        '''
    else:
        fun_text += '''
    u64 time_end=bpf_ktime_get_ns();
    struct perf_delta perf_data = {} ;
    u64 *kptr = data.lookup(&time);
    if (kptr) {
        perf_data.time_delta = time_end - *kptr;
        data.update(&time, &timeid);
    } else {
        data.insert(&time, &timeid);
        return;
    }
    output.perf_submit(ctx, &perf_data, sizeof(struct perf_delta));
}
        '''
    return fun_text


tfun_bpf_fun_start = ''
tfun_bpf_fun_end = ''


class myfun:
    cpu_cycles = False
    instructions = False
    cache_hit = False
    branch_hit = False
    bus_cycles = False

    bpf_text = base_bpf
    bpf = object
    fname = ''
    sample_period = 100
    fout = sys.stdout  # 输出位置
    isok = True

    def __init__(self, fname, cpu_cycles=False, cache_hit=False,
                 branch_hit=False, bus_cycles=False):
        self.cpu_cycles = cpu_cycles
        self.cache_hit = cache_hit
        self.branch_hit = branch_hit
        self.bus_cycles = bus_cycles
        self.fname = fname

    def load(self):
        self.bpf_text += tfun_bpf_fun_start+tfun_bpf_fun_end
        if self.cpu_cycles:
            self.bpf_text += tfun_bpf_cpu_cycles
            self.bpf_text += tfun_bpf_instructions
        if self.cache_hit:
            self.bpf_text += tfun_bpf_cache_misses
            self.bpf_text += tfun_bpf_cache_references
        if self.branch_hit:
            self.bpf_text += tfun_bpf_branch_misses
            self.bpf_text += tfun_bpf_branch_instructions
        if self.bus_cycles:
            self.bpf_text += tfun_bpf_bus_cycles
        self.bpf = BPF(text=self.bpf_text)

    def attach(self):
        self.bpf.attach_kprobe(event=self.fname, fn_name=tfun_bpf_fname_start)
        self.bpf.attach_kretprobe(event=self.fname, fn_name=tfun_bpf_fname_end)
        if self.cache_hit:  # 缓存命中
            self.bpf.attach_perf_event(ev_type=PerfType.HARDWARE, ev_config=PerfHWConfig.CACHE_MISSES,
                                       fn_name=tfun_bpf_fname_cache_misses, sample_period=self.sample_period)
            self.bpf.attach_perf_event(ev_type=PerfType.HARDWARE, ev_config=PerfHWConfig.CACHE_REFERENCES,
                                       fn_name=tfun_bpf_fname_cache_references, sample_period=self.sample_period)
        if self.branch_hit:  # 分支预测
            self.bpf.attach_perf_event(ev_type=PerfType.HARDWARE, ev_config=PerfHWConfig.BRANCH_MISSES,
                                       fn_name=tfun_bpf_fname_branch_misses, sample_period=self.sample_period)
            self.bpf.attach_perf_event(ev_type=PerfType.HARDWARE, ev_config=PerfHWConfig.BRANCH_INSTRUCTIONS,
                                       fn_name=tfun_bpf_fname_branch_instructions, sample_period=self.sample_period)
        if self.cpu_cycles:  # CPU CYCLES
            self.bpf.attach_perf_event(ev_type=PerfType.HARDWARE, ev_config=PerfHWConfig.CPU_CYCLES,
                                       fn_name=tfun_bpf_fname_cpu_cycles, sample_period=self.sample_period)
            self.bpf.attach_perf_event(ev_type=PerfType.HARDWARE, ev_config=PerfHWConfig.INSTRUCTIONS,
                                       fn_name=tfun_bpf_fname_instructions, sample_period=self.sample_period)
        if self.bus_cycles:  # BUS CYCLES
            self.bpf.attach_perf_event(ev_type=PerfType.HARDWARE, ev_config=PerfHWConfig.BUS_CYCLES,
                                       fn_name=tfun_bpf_fname_bus_cycles, sample_period=self.sample_period)

    def print_output(self, cpu, data, size):
        output = self.bpf[fun_output].event(data)
        print('\t%ld' % (output.time_delta), file=self.fout)

    def sigint_handler(self, signum, frame):
        signal.alarm(0)
        self.isok = False

    def run(self, timenum=0, file=''):
        self.bpf[fun_output].open_perf_buffer(self.print_output)
        if file != '':
            try:
                self.fout = open(file, 'w')
            except Exception as err:
                print(err)
                print('文件%s打开失败' % file)
                exit(0)
        print(self.fname, file=self.fout)
        if timenum != 0:
            print('执行%d秒，可按Ctrl^C终止' % timenum)
            signal.signal(signal.SIGALRM, self.sigint_handler)
            signal.alarm(timenum)
        else:
            print('按Ctrl^C终止')
        print('\ttime', file=self.fout)
        try:
            while self.isok:
                self.bpf.perf_buffer_poll()
        except KeyboardInterrupt:
            print()
        print(file=self.fout, flush=True)
        # 获取数据
        try:
            if self.cache_hit:
                cache_miss_count = self.bpf[tfun_bpf_hash_name][c_uint(
                    tfun_bpf_hash.CACHE_MISSES)].value
                cache_ref_count = self.bpf[tfun_bpf_hash_name][c_uint(
                    tfun_bpf_hash.CACHE_REFERENCES)].value
        except Exception:
            print('获取缓存命中数据失败', file=self.fout)
            self.cache_hit = False
        try:
            if self.branch_hit:
                branch_miss_count = self.bpf[tfun_bpf_hash_name][c_uint(
                    tfun_bpf_hash.BRANCH_MISSES)].value
                branch_inst_count = self.bpf[tfun_bpf_hash_name][c_uint(
                    tfun_bpf_hash.BRANCH_INSTRUCTIONS)].value
        except Exception:
            print('获取分支数据失败', file=self.fout)
            self.branch_hit = False
        try:
            if self.cpu_cycles:
                cpu_cycles_count = self.bpf[tfun_bpf_hash_name][c_uint(
                    tfun_bpf_hash.CPU_CYCLES)].value
                instructions_count = self.bpf[tfun_bpf_hash_name][c_uint(
                    tfun_bpf_hash.INSTRUCTIONS)].value
        except Exception:
            print('获取CPU CYCLE失败', file=self.fout)
            self.cpu_cycles = False
        try:
            if self.bus_cycles:
                bus_cycles_count = self.bpf[tfun_bpf_hash_name][c_uint(
                    tfun_bpf_hash.BUS_CYCLES)].value
        except Exception:
            print('获取BUS CYCLES失败')
            self.bus_cycles = False
        # 清理bpf
        self.bpf.cleanup()
        # 计算缓存命中率
        if self.cache_hit:
            cache_hit = (
                cache_ref_count-cache_miss_count) if (cache_ref_count >= cache_miss_count) else 0
            print('缓存命中：%ld' % cache_miss_count, file=self.fout)
            print('缓存引用：%ld' % cache_ref_count, file=self.fout)
            print('缓存命中率：%f' % (
                (float(cache_hit)/float(cache_ref_count))*100.0), file=self.fout)
        # 计算分支命中率
        print(file=self.fout)
        if self.branch_hit:
            print('分支命中：%ld' % branch_miss_count, file=self.fout)
            print('分支指令：%ld' % branch_inst_count, file=self.fout)
            print('分支预测命中率：%f' % ((float(branch_miss_count) /
                                   float(branch_inst_count))*100.0), file=self.fout)
        # 打印cpu周期
        print(file=self.fout)
        if self.cpu_cycles:
            print('CPU CYCLES: %ld' % cpu_cycles_count, file=self.fout)
            print('INSTRUCTIONS: %ld' % instructions_count, file=self.fout)
        # 打印总线周期
        if self.bus_cycles:
            print('BUS CYCLES: %ld' % bus_cycles_count, file=self.fout)


def main():
    global tfun_bpf_fun_start, tfun_bpf_fun_end
    parser = argparse.ArgumentParser(description='hardware events探测程序')
    parser.add_argument('-e', '--event', type=str,
                        default='', help='要跟踪的内核函数')
    parser.add_argument('-t', '--time', type=int,
                        default=0, help='指定探测时间，单位：秒')
    parser.add_argument('-o', '--output', type=str,
                        default='', help='指定保存数据的文件')
    parser.add_argument('-c', '--sample_period', type=int, default=100,
                        help='设置hardware事件的采样周期')
    parser.add_argument('--cache', action='store_true', help='缓存命中率')
    parser.add_argument('--branch', action='store_true', help='分支预测命中率')
    parser.add_argument('--cycle', action='store_true', help='cpu周期')
    parser.add_argument('--buscycle', action='store_true', help='总线周期')
    parser.add_argument('--print', action='store_true', help='打印bpf代码')
    parser.add_argument('--distime', action='store_true', help='阻止探测时间')

    args = parser.parse_args()
    if args.event == '':
        return
    tfun_bpf_fun_start = get_tfun_bpf_start(args.distime)
    tfun_bpf_fun_end = get_tfun_bpf_end(args.distime)

    fun = myfun(args.event)
    fun.sample_period = args.sample_period
    if args.cache:
        fun.cache_hit = True
    if args.branch:
        fun.branch_hit = True
    if args.cycle:
        fun.cpu_cycles = True
    if args.buscycle:
        fun.bus_cycles = True
    fun.load()
    if args.print:
        print(fun.bpf_text)
    fun.attach()
    fun.run(args.time, args.output)


if __name__ == '__main__':
    main()
